import json
import multiprocessing as mp
import os
import threading
import time
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from http.server import HTTPServer, BaseHTTPRequestHandler

import requests
from tqdm import tqdm


class ProgressHandler(BaseHTTPRequestHandler):
    def __init__(self, request, client_address, server, progress_bar, lock):
        self.lock = lock
        self.progress_bar = progress_bar
        super().__init__(request, client_address, server)

    def log_message(self, format, *args):
        pass

    def do_POST(self):
        if self.path == '/update':
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)
            data = json.loads(post_data.decode())

            if 'update' in data and isinstance(data['update'], int):
                with self.lock:
                    self.progress_bar.update(data['update'])

            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            self.wfile.write(json.dumps({"status": "ok"}).encode())
        else:
            self.send_response(404)
            self.end_headers()

class ProgressServer:
    port = 19674

    def __init__(self, total=0, desc=''):
        self._self_server_thread = None
        self._self_server: HTTPServer | None = None
        self.lock = threading.Lock()
        self.progress_bar = tqdm(total=total, desc=desc)

    def run_server(self):
        handler = partial(ProgressHandler, progress_bar=self.progress_bar, lock=self.lock)
        """运行HTTP服务器"""
        server = HTTPServer(('localhost', ProgressServer.port), handler)
        server_thread = threading.Thread(target=server.serve_forever)
        server_thread.daemon = True  # 设置为守护线程，这样主程序退出时服务器也会退出
        server_thread.start()
        self._self_server = server
        self._self_server_thread = server_thread

    @staticmethod
    def update_progress(n=1):
        """增加进度计数"""
        done = False
        max_try = 5
        while not done and max_try > 0:
            try:
                requests.post(f'http://localhost:{ProgressServer.port}/update', json={"update": n})
                done = True
            except requests.exceptions.RequestException:
                time.sleep(0.5)
                max_try -= 1


    def stop_server(self):
        self._self_server.shutdown()
        self._self_server.server_close()
        self._self_server_thread.join()


def process_sub_task(i):
    try:
        # 模拟实际处理过程
        time.sleep(0.2)
        return i
    except Exception as e:
        print(f"Error in process {os.getpid()}: {str(e)}")
        return False
    finally:
        ProgressServer.update_progress()

def process_task(inner_task):
    """实际的任务处理函数"""
    try:
        # 模拟实际处理过程
        time.sleep(0.1)

        tasks = [i for i in range(inner_task)]

        with ProcessPoolExecutor(max_workers=4) as executor:
            result = list(executor.map(process_sub_task, tasks))

        return result
    except Exception as e:
        print(f"Error in process {os.getpid()}: {str(e)}")
        return False

def run_parallel_tasks(outer_total, inner_total, num_processes=None):
    """运行并行任务的主函数"""
    if num_processes is None:
        num_processes = mp.cpu_count()

    # 计算总任务数并初始化进度
    total_tasks = outer_total * inner_total

    s = ProgressServer(total_tasks, 'Progressing')
    s.run_server()

    # 准备任务参数
    tasks = [inner_total for _ in range(outer_total)]

    with ProcessPoolExecutor(max_workers=num_processes) as executor:
        # 返回(git, {module: unit_map})
        result = list(executor.map(process_task, tasks))

    # 关闭服务器
    s.stop_server()

    print(result)

def main():
    # 设置参数
    outer_total = 5    # 外层循环总数
    inner_total = 10   # 内层循环总数
    num_processes = 4  # 进程数
    port = 8000        # 服务器端口

    # 运行任务
    run_parallel_tasks(
        outer_total=outer_total,
        inner_total=inner_total,
        num_processes=num_processes
    )

if __name__ == '__main__':
    main()